import pandas as pd
import torch
from transformers import BertTokenizer

import config


class ContradictionDataset(torch.utils.data.Dataset):
    def __init__(self, csv_file_path: str):
        self.tokenizer = BertTokenizer.from_pretrained(
            config.bert_pretrained_folder)
        self.texts, self.labels = self.read_csv_file(csv_file_path)

    def __getitem__(self, idx):
        encoding = self.tokenizer(self.texts[idx],
                                  truncation=True,
                                  max_length=512)
        item = {k: torch.tensor(v) for k, v in encoding.items()}
        item["labels"] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)

    @staticmethod
    def read_csv_file(csv_file_path: str):
        data = pd.read_csv(csv_file_path)
        assert "content" in data.columns.tolist()
        assert "category" in data.columns.tolist()
        # 删除内容为空的行
        data = data.dropna(axis=0)
        # category列中的类别是汉字，根据类别的数量，转换成对应的类别
        texts = data["content"].tolist()
        labels = [config.categories.index(c) for c in data["category"]]
        assert (all([label >= 0 for label in labels]))
        return texts, labels


def main():
    dataset = ContradictionDataset(config.train_csv_file_path)
    print(dataset[0]["input_ids"].shape)
    print(dataset[1]["input_ids"].shape)


if __name__ == "__main__":
    main()
